package com.wzw.springboot.principle.use.websock.second;

import com.wzw.springboot.principle.use.json.JacksonHandler;
import com.wzw.springboot.principle.use.security.infrastructure.model.SecurityUser;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Wangzhiwen
 */
@Slf4j
public class WebSocketSessionManager {

    private static final String KEY = SecurityUser.class.getSimpleName();

    private static final Map<Long, WebSocketSession> WEB_SOCKET_SESSION_MAP = new ConcurrentHashMap<>(16);

    public static <T> void sendMessage(WebSocketSession session, WebSocketMessage<T> message) {
        try {
            session.sendMessage(message);
        } catch (IOException e) {
            log.error("WenSocket Send Message Error", e);
        }
    }

    public static <T> void sendMessage(Long userId, WebSocketMessage<T> message) {
        WebSocketSession session = WEB_SOCKET_SESSION_MAP.get(userId);
        sendMessage(session, message);
    }

    public static void sendJsonMessage(Long userId, Object message) {
        WebSocketSession session = WEB_SOCKET_SESSION_MAP.get(userId);
        String jsonString = JacksonHandler.toJsonString(message);
        sendMessage(session, new TextMessage(jsonString));
    }

    public static void addSession(WebSocketSession session) {
        SecurityUser securityUser = (SecurityUser) session.getAttributes().get(KEY);
        WEB_SOCKET_SESSION_MAP.put(securityUser.getId(), session);
    }

    public static void removeSession(WebSocketSession session) {
        SecurityUser securityUser = (SecurityUser) session.getAttributes().get(KEY);
        WEB_SOCKET_SESSION_MAP.remove(securityUser.getId());
    }
}
